import torch
import librosa
import numpy as np

# --- PyTorch/CUDA 设备设置 ---
# 如果有 CUDA 可用，使用 GPU，否则使用 CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def preemphasis(wav, k, device=None):
    """预加重滤波"""
    if not isinstance(wav, torch.Tensor):
        wav = torch.tensor(wav, dtype=torch.float32)
    if device:
        wav = wav.to(device)
    return torch.cat([wav[0:1], wav[1:] - k * wav[:-1]])

def _stft(y, device=None):
    """短时傅里叶变换"""
    if not isinstance(y, torch.Tensor):
        y = torch.tensor(y, dtype=torch.float32)
    if device:
        y = y.to(device)
    
    # 使用PyTorch的STFT
    return torch.stft(
        y, 
        n_fft=800, 
        hop_length=200, 
        win_length=800,
        window=torch.hann_window(800, device=y.device),
        return_complex=True
    )

def _build_mel_basis(device=None):
    """构建梅尔滤波器组"""
    mel_basis = librosa.filters.mel(
        sr=16000,
        n_fft=800,
        n_mels=80,
        fmin=55,
        fmax=7600
    )
    
    mel_basis = torch.tensor(mel_basis, dtype=torch.float32)
    if device:
        mel_basis = mel_basis.to(device)
    
    return mel_basis

def _linear_to_mel(spectrogram, mel_basis=None, device=None):
    """线性频谱转梅尔频谱"""
    if device is None:
        device = spectrogram.device
    if not isinstance(spectrogram, torch.Tensor):
        spectrogram = torch.tensor(spectrogram, dtype=torch.float32)
    
    if mel_basis is None:
        mel_basis = _build_mel_basis(device)
    
    return torch.matmul(mel_basis, spectrogram)

def _amp_to_db(x, device=None):
    """幅度转分贝"""
    if device:
        x = x.to(device)
    
    min_level = torch.tensor(np.exp(-5 * np.log(10)), dtype=torch.float32, device=x.device)
    return 20 * torch.log10(torch.clamp(x, min=min_level))

def _normalize(S):
    """标准化处理"""
    # 原始代码: return np.clip((2 * 4.) * ((S - -100) / (--100)) - 4., -4., 4.)
    # 注意: (--100) 实际上是 100，因为双负号变正号
    return torch.clamp((2 * 4.) * ((S - (-100)) / 100) - 4., -4., 4.)

def melspectrogram(wav, device=None):
    """计算梅尔频谱图"""
    if not isinstance(wav, torch.Tensor):
        wav = torch.tensor(wav, dtype=torch.float32)
    if device:
        wav = wav.to(device)
    # 预加重处理
    wav_emph = preemphasis(wav, 0.97, device)
    
    # STFT
    D = _stft(wav_emph, device)
    
    # 计算幅度
    magnitude = torch.abs(D)
    
    # 转换为梅尔频谱
    mel_basis = _build_mel_basis(device)
    mel_spec = _linear_to_mel(magnitude, mel_basis, device)
    
    # 转换为分贝单位
    S = _amp_to_db(mel_spec, device) - 20
    
    # 标准化
    normalized = _normalize(S)
    
    return normalized